package cas.ihep.hep.util;

import cas.ihep.fs.UnifiedFileSystem;
import cas.ihep.hep.analysis.NewLoopPWA;
import cas.ihep.spark.io.Transformer;

import java.io.*;
import java.net.URI;
import java.net.URISyntaxException;
import java.nio.ByteBuffer;
import java.util.ArrayList;
import java.util.List;
import java.util.Random;
import java.util.concurrent.atomic.AtomicInteger;

public class EventBlock implements Serializable {
    private int eventCount;
    private URI filePath;
    private EventBlock(){
    }

    public URI getPath(){
        return filePath;
    }
    private transient NativeBlock nBlock=null;

    public NativeBlock nativeData()throws IOException{
        if(nBlock==null){
            nBlock=new NativeBlock(filePath);
        }
        return nBlock;
    }

    public int count(){
        return eventCount;
    }

    private static final String BaseURI="alluxio://hadoop08.ihep.ac.cn:19998/tmp/EventBlock-";
    private static final AtomicInteger accumulator=new AtomicInteger();
    private static final Random rand=new Random();

    public static List<EventBlock> splitBlock(InputStream in,int eleSize,int eleCount,int partitions)throws IOException{
        int dteach= eleCount/partitions;
        if(eleCount%partitions!=0){
            dteach++;
        }
        int each=dteach;
        ArrayList<EventBlock> ret=new ArrayList<>();
        Transformer.Source src=Transformer.fromInputStream(in,eleSize,eleCount);
        try (Transformer.MultiSink dest = new Transformer.MultiSink() {
            int current = 0;
            Transformer.Sink out = null;

            Transformer.Sink createSink() throws IOException {
                EventBlock block = new EventBlock();
                URI uri;
                do {
                    try {
                        uri = new URI(BaseURI + Integer.toHexString(rand.nextInt()) + '-' + accumulator.getAndIncrement());
                    } catch (URISyntaxException ignored) {
                        throw new IllegalStateException("Should never happens");
                    }
                } while (NewLoopPWA.getUnifiedFileSystem().exists(uri));
                block.filePath=uri;
                ret.add(block);
                return new Transformer.OutputStreamSink(NewLoopPWA.getUnifiedFileSystem().createStream(uri),true){
                    int evtCount=0;
                    public void put(ByteBuffer buf)throws Exception{
                        super.put(buf);
                        evtCount++;
                    }
                    public void close()throws IOException{
                        block.eventCount=evtCount;
                        super.close();
                    }
                };
            }

            @Override
            public Transformer.Sink currentSink() throws Exception {
                if (current >= each) {
                    current = 0;
                    out.close();
                    out = createSink();
                } else if (out == null) {
                    out = createSink();
                }
                current++;
                return out;
            }

            @Override
            public void close() throws IOException {
                if (out != null) {
                    out.close();
                }
            }
        }) {
            Transformer.transform(src, dest);
        } catch (IOException e) {
            throw e;
        } catch (Exception e) {
            throw new IOException(e);
        }
        return ret;
    }
    public static EventBlock readFromFile(UnifiedFileSystem ufs,InputStream in,int length,int evt)throws IOException{
        EventBlock block=new EventBlock();
        URI uri;
        do{
            try {
                uri=new URI(BaseURI+Integer.toHexString(rand.nextInt())+'-'+accumulator.getAndIncrement());
            } catch (URISyntaxException ignored) {
                throw new IllegalStateException("Should never happens");
            }
        }while(ufs.exists(uri));
        try(OutputStream bos=ufs.createStream(uri)){
            byte[] buf=new byte[4096];
            int total=0;
            while(total<length) {
                int each = Math.min(4096, length - total);
                int num=in.read(buf,0,each);
                if(num!=each){
                    throw new EOFException("Unexpected EOF when transforming file");
                }
                bos.write(buf,0,each);
                total+=each;
            }
        }
        block.filePath=uri;
        block.eventCount=evt;
        return block;
    }

    public void close(){
        try {
            NewLoopPWA.getUnifiedFileSystem().delete(filePath);
        } catch (IOException ignored) {
        }
    }
}
